#!/bin/bash

########### Run enviroment check
rpm -qa | grep expect-[0-9] >>/dev/null
rc=$?
if [[ $rc -ne 0 ]];then
    echo "Package expect is not installed,please install it first !";
    exit 10;
fi

function echo_error(){
    echo -e "\033[31;1m[ERROR] $1 \033[0m";
}

function echo_msg() {
    echo -e "\033[34;1m[ MSG ] $1 \033[0m";
}

BASE_PATH=`pwd`
LOCAL_HOST=192.168.23.1
RESULT_FILE_PATH=${BASE_PATH}/result
HOST_CONF_FILE=${BASE_PATH}/conf/host.conf
TEST_CONF_FILE=${BASE_PATH}/conf/init.conf
FIO_TOOL=${BASE_PATH}/tool/fio
REMOTE_BASE_PATH=/home/perf_test
REMOTE_TEST_PATH=${REMOTE_BASE_PATH}/perf_test
# Default interval time between each fio job
JOB_INTERVAL_TIME=60
# Total fio job number
JOB_NUM=0
# Total host number
HOST_NUM=0

function base_file_check() {
    if [[ ! -d ${RESULT_FILE_PATH} ]];then
        mkdir ${RESULT_FILE_PATH}
    else
        ls ${RESULT_FILE_PATH} | grep result\\.log >> /dev/null 2>&1
        rc=$?
        if [[ rc -eq 0 ]];then
	    bak=${RESULT_FILE_PATH}/`date +%F_%T`
            mkdir ${bak}
            mv ${RESULT_FILE_PATH}/*.log ${bak}
        fi
    fi

    if [[ ! -f ${FIO_TOOL} ]] ;then
        echo_error "The fio tool ./tool/fio is not exist!"
        exit 10
    else
        chmod 755 ${FIO_TOOL}
    fi

    if [[ ! -f ${HOST_CONF_FILE} ]]; then
        echo_error "Configuration file host.conf is not exist!"
        exit 10
    fi

    if [[ ! -f ${TEST_CONF_FILE} ]]; then
        echo_error "Test configuration file init.conf is not exist!"
        exit 10
    fi
}

declare -a host
declare -a password
declare -a device
declare -a blocksize
declare -a rw_model
declare -a runtime
declare -a iodepth
declare -a read_ratio

function init_data(){
    echo "[INFO] Data initializing..."
    
    JOB_INTERVAL_TIME=`sed -n '/^#/d;/JOB_INTERVAL/p' ${TEST_CONF_FILE}  | awk -F = '{ if ($2 ~ /^[0-9]+$/ ) {print $2;} else {print "NULL";}}'`
    if [[ ${JOB_INTERVAL_TIME} == "NULL" ]]; then
        echo_error "Parameter JOB_INTERVAL is not correct,please check !";
        exit 10;
    fi
    
    host=(`sed '/^ *$/d;/^#/d' ${HOST_CONF_FILE}  | grep -v "\[.*\]" | awk '{ if (NF == 3){ print $1;} }' `)
    for i in ${host[@]};do
        if [[ ${LOCAL_HOST} == "0.0.0.0" ]];then
            ifconfig | grep $i > /dev/null
            rc=$?
            if [ $rc -eq 0 ];then
                LOCAL_HOST=$i;
                echo "[INFO] Local host is :${LOCAL_HOST}.";
            fi
        fi        
        if [[ $i =~ ^[1-9][0-9]{0,2}(\.[0-9]{1,3}){3}$ ]];then
            continue;
        else 
            echo_error "Host ip $i is not valid,please check file host.conf !";
            exit 10;            
        fi        
    done
    HOST_NUM=${#host[@]}
    
    if [[ ${HOST_NUM} -eq 0 ]]; then
        echo_error "No valid hosts,please check file host.conf !";
        exit 10;
    fi
    
    password=(`sed '/^ *$/d;/^#/d' ${HOST_CONF_FILE}  | grep -v "\[.*\]" | awk '{ if (NF == 3){ print $2;} }' `)
	
    device=(`sed '/^ *$/d;/^#/d' ${HOST_CONF_FILE}  | grep -v "\[.*\]" | awk '{ if (NF == 3){ print $3;} }' `)
    
    blocksize=(`sed -n "/^ *$/d;/^#/d;/\[FIO_JOB\]/,/\[.*\]/p" ${TEST_CONF_FILE}  | grep -v "\[.*\]" | awk '{ if ((NF == 4)||(NF == 5)) { print $1;} }' `)
    for i in ${blocksize[@]};do
        if [[ $i =~ ^[0-9]+k(-[0-9]+k)?$ ]];then
            continue;
        else 
            echo_error "Parameter blocksize is not valid,please check file init.conf !";
            exit 10;            
        fi    
    done
    JOB_NUM=${#blocksize[@]}
    
    if [[ ${JOB_NUM} -eq 0 ]]; then
        echo_error "No valid fio jobs,please check file init.conf !";
        exit 10;
    fi
    rw_model=(`sed -n "/^ *$/d;/^#/d;/\[FIO_JOB\]/,/\[.*\]/p" ${TEST_CONF_FILE}  | grep -v "\[.*\]" | awk '{ if ((NF == 4)||(NF == 5)) { print $2;} }' `)

    runtime=(`sed -n "/^ *$/d;/^#/d;/\[FIO_JOB\]/,/\[.*\]/p" ${TEST_CONF_FILE}  | grep -v "\[.*\]" | awk '{ if ((NF == 4)||(NF == 5)) { print $3;} }' `)
    for i in ${runtime[@]};do
        if [[ $i =~ ^[0-9]+$ ]];then
            continue;
        else 
            echo_error "Parameter runtime is not valid,please check file init.conf !";
            exit 10;            
        fi    
    done
	
    iodepth=(`sed -n "/^ *$/d;/^#/d;/\[FIO_JOB\]/,/\[.*\]/p" ${TEST_CONF_FILE}  | grep -v "\[.*\]" | awk '{ if ((NF == 4)||(NF == 5)) { print $4;} }' `)
    for i in ${iodepth[@]};do
        if [[ $i =~ ^[0-9]+$ ]];then
            if [[ ( $i -gt 128 ) ]]; then 
                echo_error "Parameter iodepth is bigger than 128,please adjust it !";
                exit 10;
            fi
        else 
            echo_error "Parameter iodepth is not valid,please check file init.conf !";
            exit 10;            
        fi
    done
    read_ratio=(`sed -n "/^ *$/d;/^#/d;/\[FIO_JOB\]/,/\[.*\]/p" ${TEST_CONF_FILE}  | grep -v "\[.*\]" | awk '{ if ((NF == 4)||(NF == 5)) { if ( $5 == "" ) { print "NULL";} else print $5;} }' `)
    for i in ${read_ratio[@]};do
        if [[ ($i =~ ^[0-9]+$) && ( $i -gt 0 ) && ( $i -lt 100 ) ]];then
            continue;
        else if [[ $i != "NULL" ]];then
            echo_error "Parameter read_ratio is not valid,please check file init.conf !";
            exit 10;
            fi
        fi
    done    
    
    echo_msg "Total valid hosts num:${HOST_NUM} !"
    echo_msg "Total valid jobs num:${JOB_NUM} !"
    echo "[INFO] Done,data initialized!"     
}

function generate_test_conf(){
    echo "[INFO] Begin to generate fio configuration file !"    
    seq=0
    while [[ seq -lt ${JOB_NUM} ]]
    do
        TEMP_CONF_FILE="${BASE_PATH}/conf/FIO_JOB_${seq}.fio"
        sed -n '/^#/d;/^ *$/d;/\[global\]/,/\[*\]/p' conf/init.conf | grep -v "\[[^(global)].*\]" > ${TEMP_CONF_FILE}
#       echo "filename=DEVICE"  >> ${TEMP_CONF_FILE}
        echo "[JOB_${seq}_${blocksize[seq]}_${rw_model[seq]}_DISK${seq}]"  >> ${TEMP_CONF_FILE}
#       echo "stonewall"              >> ${TEMP_CONF_FILE}
        if [[ ${blocksize[seq]} == *-* ]];then
            echo "bsrange=${blocksize[seq]}"     >> ${TEMP_CONF_FILE}
        else    
            echo "bs=${blocksize[seq]}"          >> ${TEMP_CONF_FILE}
        fi
        echo "rw=${rw_model[seq]}"     >> ${TEMP_CONF_FILE}
        if [[ ${rw_model[seq]} == *rw* ]];then
            if [[ ${read_ratio[seq]} == "NULL" ]];then
                echo "rwmixread=70"     >> ${TEMP_CONF_FILE}
                echo_msg "The read_ratio for job ${seq} is not correct or not setted,use default 70% !"
            else 
                echo "rwmixread=${read_ratio[seq]}"     >> ${TEMP_CONF_FILE}
            fi
        fi
#        echo "runtime=${runtime[seq]}" >> ${TEMP_CONF_FILE}
        echo "iodepth=${iodepth[seq]}" >> ${TEMP_CONF_FILE}
        echo "rbdname=DEVICE1"  >> ${TEMP_CONF_FILE}
	
	((seq++))
    done
    echo "[INFO] Generate fio test configuration file succeed !"    
}

## 
## $1:host_ip        $2:root password
##
function generate_pub_key(){    
    expect -c "set timeout 3;
        spawn ssh root@$1 ssh-keygen -t rsa
        expect {
            -re \"(yes/no)\" {send \"yes\r\";exp_continue;}
            -re \"assword:\" {send \"$2\r\";}
        };
        expect {
            -re \"assword:\" {exit 98;}
            -re \"Enter\"    {send \"\r\";exp_continue}
            -re \"(y/n)\"    {send \"y\r\";exp_continue}
            timeout          {exit 99;}
            eof              {exit 0;}
        }" >> /dev/null;
               
    result=$?
    if [[ ${result} -eq 0 ]];then
#        echo "[INFO] Generate public file on host $1 succeed!"
        return 0;
    else if [[ ${result} -eq 98 ]]; then
        echo_error "Password for $1 is not correct !"
        exit 98;
        fi
        if [[ ${result} -eq 99 ]]; then
        echo "[WARN] Server $1 response timeout!"
        exit 99;
        fi
    fi
}

## 
## $1:source_file        $2:dest-file
## $3:host_ip            $4:root password
##
function copy_file(){
    expect -c "set timeout 5;
        spawn scp $1 $2 
        expect {
            -re \"(yes/no)\" {send \"yes\r\";exp_continue;}
            -re \"assword:\" {send \"$4\r\";}
            eof              {exit 0;}
        };
        expect {
            -re \"assword:\" {exit 98;}
            timeout          {exit 99;}
            eof              {exit 0;}
        }" >> /dev/null ;
    result=$?
    if [[ ${result} -eq 98 ]]; then
        echo_error "Password for host $3 is not correct !"
        exit 98;
    fi
    if [[ ${result} -eq 99 ]]; then
        echo "[WARN] Server $3 response timeout!"
        exit 99;
    fi
    if [[ ${result} -eq 0 ]]; then
        return 0;
    fi
}

function make_trust(){
    echo "[INFO] Begin to config ssh trust between each host !"
    if [ -f ${RESULT_FILE_PATH}/authorized_keys ];then
        rm -rf ${RESULT_FILE_PATH}/authorized_keys    
    fi
    count=0
    for h in ${host[@]}
    do
        generate_pub_key $h ${password[count]};
        copy_file root@$h:/root/.ssh/id_rsa.pub ${RESULT_FILE_PATH}/${h}.pub $h ${password[count]}
        echo "[INFO] Copy file id_rsa.pub from $h succeed !"
        ((count++))
    done
    cat ${RESULT_FILE_PATH}/*.pub > ${RESULT_FILE_PATH}/authorized_keys
    count=0
    for h in ${host[@]}
    do
        copy_file ${RESULT_FILE_PATH}/authorized_keys root@$h:/root/.ssh/authorized_keys $h ${password[count]}
        echo "[INFO] Copy file authorized_keys to $h succeed !"
        ((count++))
    done
    echo "[INFO] Config ssh trust between each host succeed!"
    rm -rf ${RESULT_FILE_PATH}/*.pub
    rm -rf ${RESULT_FILE_PATH}/authorized_keys
}

##
## $1 host_ip
##
function check_run_env(){
    ssh $1 'rpm -qa | grep libaio-[0-9] >/dev/null 2>&1'
    rc1=$?
    if [ $rc1 -ne 0 ];then 
        echo_error "The package libaio is not installed on $h ,exit !";
        exit 95;    
    fi
#    ssh $1 'rpm -qa | grep libaio-devel >/dev/null 2>&1'
#    rc2=$?
#    if [ $rc2 -ne 0 ];then 
#        echo_error "The package libaio-devel is not installed on $h ,exit !";
#        exit 95;
#    fi
    ssh $1 'rpm -qa | grep expect-[0-9] >/dev/null 2>&1'
    rc3=$?
    if [ $rc3 -ne 0 ];then
        echo_error "The package expect is not installed on $h ,exit !";
        exit 95;
    fi
}

##
## $1:scripts file name 
##
function generate_stop_scripts() {
    echo "#!/bin/bash"            > $1
    echo "for proc in \`ps -ef | grep fio | grep -v grep | awk '{ print \$2 }'\`" >> $1
    echo "do "                    >> $1
    echo '    kill -9 ${proc}'    >> $1
    echo "done"                   >> $1
}

##
## $1:scripts file name 
##
function generate_fio_scripts() {
    echo "#!/bin/bash"                      > $1
    echo "job_num=${JOB_NUM}"               >> $1
    echo "base_path=${REMOTE_TEST_PATH}"    >> $1
    echo "result_file=\${base_path}/LOCAL_IP_result.log" >>$1
    echo "i=0"                                           >> $1
    echo 'for f in `ls ${base_path}/FIO_JOB_*`'          >> $1
    echo "do"                                            >> $1
    echo '    iostat -xm 10 >> ${base_path}/JOB_${i}_iostat.log & '     >> $1
    echo '    mpstat -P ALL 10 >> ${base_path}/JOB_${i}_mpstat.log & '  >> $1
    echo '    ${base_path}/fio --output=${base_path}/JOB_${i}.log $f '  >> $1
    echo "    sleep 3"               >> $1
    echo "    killall -9 iostat "    >> $1
    echo "    killall -9 mpstat "    >> $1
    echo "    cat \${base_path}/JOB_\${i}.log | egrep '(JOB.*group)|(iops=)|(^\s*lat.*avg)' >> \${result_file}" >> $1
    echo '    ((i++))'                           >> $1
    echo '    if [ ${i} -lt ${job_num} ]; then'  >> $1
    echo "        sleep ${JOB_INTERVAL_TIME}"    >> $1
    echo "    fi"                    >> $1
    echo "done"                      >> $1    
    echo 'expect -c "set timeout 3;'             >> $1
    echo "spawn scp \${result_file} root@${LOCAL_HOST}:${RESULT_FILE_PATH}/" >> $1
    echo "expect {"                  >> $1
    echo '    -re \"(yes/no)\" {send \"yes\r\";exp_continue;}' >> $1
    echo '    timeout          {exit 99;}'  >> $1
    echo '    eof              {exit 0;}'   >> $1
    echo '};" >/dev/null '                  >> $1
}

function copy_tool_and_conf(){
    
    stop_test=${BASE_PATH}/stop_test.sh
    start_test=${BASE_PATH}/start_fio_test_tmp.sh
    start_fio_test=${BASE_PATH}/start_fio_test.sh

    generate_stop_scripts ${stop_test}
    generate_fio_scripts  ${start_test}
    
    count=0
    for h in ${host[@]}
    do
        check_run_env $h
        ssh $h ls ${REMOTE_TEST_PATH} > /dev/null  2>&1
        rc=$?
        if [ $rc -ne 0 ];then 
            ssh $h "mkdir -p ${REMOTE_TEST_PATH}"
        else
            ssh $h "rm -rf ${REMOTE_TEST_PATH}/*"
        fi
        for f in `ls ${BASE_PATH}/conf/FIO_JOB_*`
        do
	    DISK1=`echo ${device[count]} | awk -F "," '{print $1}'`
            DISK2=`echo ${device[count]} | awk -F "," '{print $2}'`

	    sed "s#DEVICE1#${DISK1}#g;s#DEVICE2#${DISK2}#g" $f > ${f}_${h}
#	    sed s#DEVICE2#${DISK2}#g $f > ${f}_${h}
            scp ${f}_${h} root@${h}:${REMOTE_TEST_PATH} > /dev/null  2>&1
            rc=$?
            if [ $rc -eq 0 ];then
                echo "[INFO] Copy conf file ${f} to $h succeed ! "
            else
               echo_error " Copy conf file ${f} to $h failed !";
                exit 96;
            fi
            rm -rf ${f}_${h}
        done
        sed s/LOCAL_IP/${h}/g  ${start_test} > ${start_fio_test}
        scp ${FIO_TOOL} ${stop_test} ${start_fio_test} root@${h}:${REMOTE_TEST_PATH} > /dev/null  2>&1
        if [ $rc -eq 0 ];then
            echo "[INFO] Copy scripts to $h succeed ! "
        else
            echo_error "Copy scripts to $h failed !";
            exit 96;
        fi
        ((count++))
    done
    rm -rf ${stop_test}
    rm -rf ${start_test}
    rm -rf ${start_fio_test}
    rm -rf ${BASE_PATH}/conf/FIO_JOB_*
}

function run_fio_test(){
    count=0
    for h in ${host[@]}
    do
        ssh $h "sh ${REMOTE_TEST_PATH}/stop_test.sh > /dev/null 2>&1 & "
    done
    echo "[INFO] Preparing for fio jobs......"
    sleep 5
    for h in ${host[@]}
    do
        ssh $h "sh ${REMOTE_TEST_PATH}/start_fio_test.sh > /dev/null 2>&1 &"
        rc=$?
        if [ $rc -eq 0 ];then
            echo "[INFO] Startup fio on $h succeed ! "
        else
            echo_error "Startup fio on $h failed !";
            exit 95;
        fi
        ((count++))
    done
}

function main(){
    base_file_check
    init_data
    generate_test_conf    
    make_trust
    copy_tool_and_conf    
    run_fio_test
}
main

